#include <immintrin.h>
#include <pthread.h>
#include <windows.h>
#include <iostream>
#include <random>
#include "threadpool.c"
using namespace std;
#define NUM_THREADS 4
#define N 5

float matrix[N][N];
pthread_t thread[NUM_THREADS];
pthread_mutex_t lock;
int done;

void createMatrix() {
    srand(20000214);
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < N; j++) {
            matrix[i][j] = rand() % 100;
            cout << matrix[i][j] << " ";
        }
        cout << endl;
    }
    cout << endl;
}

typedef struct {
    int threadId;
    int k;
} threadParm_t;

threadParm_t parm[NUM_THREADS];

pthread_barrier_t barrier;

void LU() {
    for (int k = 0; k < N; k++) {
        for (int j = k + 1; j < N; j++) {
            matrix[k][j] /= matrix[k][k];
        }
        matrix[k][k] = 1;
        for (int i = k + 1; i < N; i++) {
            for (int j = k + 1; j < N; j++) {
                matrix[i][j] -= matrix[i][k] * matrix[k][j];
            }
            matrix[i][k] = 0;
        }
    }
}

void* pthread_col_function(void* arg) {
    threadParm_t* parm = (threadParm_t*)arg;
    int k = parm->k;
    int len = (N - k) / 4;
    int start, end;
    start = k + parm->threadId * len;
    if (start == k) {
        start++;
    }
    if (parm->threadId != NUM_THREADS - 1) {
        end = k + (parm->threadId + 1) * len;
    } else {
        end = N;
    }
    for (int j = start; j < end; j++) {
        matrix[k][j] /= matrix[k][k];
    }
    for (int i = k + 1; i < N; i++) {
        for (int j = start; j < end; j++) {
            matrix[i][j] -= matrix[i][k] * matrix[k][j];
        }
    }
    pthread_barrier_wait(&barrier);
    return NULL;
}

void* pthread_col_pool_function(void* arg) {
    threadParm_t* parm = (threadParm_t*)arg;
    int k = parm->k;
    int len = (N - k) / 4;
    int start, end;
    start = k + parm->threadId * len;
    if (start == k) {
        start++;
    }
    if (parm->threadId != NUM_THREADS - 1) {
        end = k + (parm->threadId + 1) * len;
    } else {
        end = N;
    }
    for (int j = start; j < end; j++) {
        matrix[k][j] /= matrix[k][k];
    }
    for (int i = k + 1; i < N; i++) {
        for (int j = start; j < end; j++) {
            matrix[i][j] -= matrix[i][k] * matrix[k][j];
        }
    }
    pthread_mutex_lock(&lock);
    done++;
    pthread_mutex_unlock(&lock);
    return NULL;
}

void pthread_col() {
    for (int k = 0; k < N; k++) {
        pthread_barrier_init(&barrier, NULL, NUM_THREADS);
        for (int i = 0; i < NUM_THREADS; i++) {
            parm[i] = threadParm_t{i, k};
            pthread_create(&thread[i], NULL, pthread_col_function,
                           (void*)&parm[i]);
        }
        for (int i = 0; i < NUM_THREADS; i++) {
            pthread_join(thread[i], NULL);
        }
        pthread_barrier_destroy(&barrier);
        matrix[k][k] = 1;
        for (int i = k + 1; i < N; i++) {
            matrix[i][k] = 0;
        }
    }
}

void pthread_col_pool() {
    pthread_mutex_init(&lock, NULL);
    threadpool_t* threadpool = threadpool_create(NUM_THREADS, NUM_THREADS, 0);
    for (int k = 0; k < N; k++) {
        done = 0;
        for (int i = 0; i < NUM_THREADS; i++) {
            parm[i] = threadParm_t{i, k};
            threadpool_add(threadpool, &pthread_col_pool_function,
                           (void*)&parm[i], 0);
        }
        while (done != 4) {
            Sleep(0.01);
        }
        matrix[k][k] = 1;
        for (int i = k + 1; i < N; i++) {
            matrix[i][k] = 0;
        }
    }
    pthread_mutex_destroy(&lock);
}

void* pthread_row_function(void* arg) {
    threadParm_t* parm = (threadParm_t*)arg;
    int k = parm->k;
    int len = (N - k - 1) / NUM_THREADS;
    int start = k + 1 + len * parm->threadId, end;
    if (parm->threadId != NUM_THREADS - 1) {
        end = k + 1 + (parm->threadId + 1) * len;
    } else {
        end = N;
    }
    for (int i = start; i < end; i++) {
        for (int j = k + 1; j < N; j++) {
            matrix[i][j] -= matrix[i][k] * matrix[k][j];
        }
        matrix[i][k] = 0;
    }
    return NULL;
}

void* pthread_row_pool_function(void* arg) {
    threadParm_t* parm = (threadParm_t*)arg;
    int k = parm->k;
    int len = (N - k - 1) / NUM_THREADS;
    int start = k + 1 + len * parm->threadId, end;
    if (parm->threadId != NUM_THREADS - 1) {
        end = k + 1 + (parm->threadId + 1) * len;
    } else {
        end = N;
    }
    for (int i = start; i < end; i++) {
        for (int j = k + 1; j < N; j++) {
            matrix[i][j] -= matrix[i][k] * matrix[k][j];
        }
        matrix[i][k] = 0;
    }
    pthread_mutex_lock(&lock);
    done++;
    pthread_mutex_unlock(&lock);
    return NULL;
}

void pthread_row() {
    for (int k = 0; k < N; k++) {
        for (int j = k + 1; j < N; j++) {
            matrix[k][j] /= matrix[k][k];
        }
        matrix[k][k] = 1;
        for (int i = 0; i < NUM_THREADS; i++) {
            parm[i] = threadParm_t{i, k};
            pthread_create(&thread[i], NULL, pthread_row_function,
                           (void*)&parm[i]);
        }
        for (int i = 0; i < NUM_THREADS; i++) {
            pthread_join(thread[i], NULL);
        }
    }
}

void pthread_row_pool() {
    pthread_mutex_init(&lock, NULL);
    threadpool_t* threadpool = threadpool_create(NUM_THREADS, NUM_THREADS, 0);
    for (int k = 0; k < N; k++) {
        done = 0;
        for (int j = k + 1; j < N; j++) {
            matrix[k][j] /= matrix[k][k];
        }
        matrix[k][k] = 1;
        for (int i = 0; i < NUM_THREADS; i++) {
            parm[i] = threadParm_t{i, k};
            threadpool_add(threadpool, &pthread_row_pool_function,
                           (void*)&parm[i], 0);
        }
        while (done != 4) {
            Sleep(0.01);
        }
    }
}

void* pthread_row_sse_function(void* arg) {
    threadParm_t* parm = (threadParm_t*)arg;
    int k = parm->k;
    int len = (N - k - 1) / NUM_THREADS;
    int start = k + 1 + len * parm->threadId, end;
    if (parm->threadId != NUM_THREADS - 1) {
        end = k + 1 + (parm->threadId + 1) * len;
    } else {
        end = N;
    }
    __m128 t1, t2, t3, t4;
    for (int i = start; i < end; i++) {
        float ik[4] = {matrix[i][k], matrix[i][k], matrix[i][k], matrix[i][k]};
        t1 = _mm_loadu_ps(ik);
        for (int j = N - 4; j > k; j -= 4) {
            t2 = _mm_loadu_ps(matrix[i] + j);
            t3 = _mm_loadu_ps(matrix[k] + j);
            t4 = _mm_sub_ps(t2, _mm_mul_ps(t1, t3));
            _mm_storeu_ps(matrix[i] + j, t4);
        }
        for (int j = k + 1; j % 4 != N % 4; j++) {
            matrix[i][j] -= matrix[i][k] * matrix[k][j];
        }
        matrix[i][k] = 0;
    }
    return NULL;
}

void pthread_row_sse() {
    for (int k = 0; k < N; k++) {
        for (int j = k + 1; j < N; j++) {
            matrix[k][j] /= matrix[k][k];
        }
        matrix[k][k] = 1;
        for (int i = 0; i < NUM_THREADS; i++) {
            parm[i] = threadParm_t{i, k};
            pthread_create(&thread[i], NULL, pthread_row_sse_function,
                           (void*)&parm[i]);
        }
        for (int i = 0; i < NUM_THREADS; i++) {
            pthread_join(thread[i], NULL);
        }
    }
}

void* pthread_row_sse_pool_function(void* arg) {
    threadParm_t* parm = (threadParm_t*)arg;
    int k = parm->k;
    int len = (N - k - 1) / NUM_THREADS;
    int start = k + 1 + len * parm->threadId, end;
    if (parm->threadId != NUM_THREADS - 1) {
        end = k + 1 + (parm->threadId + 1) * len;
    } else {
        end = N;
    }
    __m128 t1, t2, t3, t4;
    for (int i = start; i < end; i++) {
        float ik[4] = {matrix[i][k], matrix[i][k], matrix[i][k], matrix[i][k]};
        t1 = _mm_loadu_ps(ik);
        for (int j = N - 4; j > k; j -= 4) {
            t2 = _mm_loadu_ps(matrix[i] + j);
            t3 = _mm_loadu_ps(matrix[k] + j);
            t4 = _mm_sub_ps(t2, _mm_mul_ps(t1, t3));
            _mm_storeu_ps(matrix[i] + j, t4);
        }
        for (int j = k + 1; j % 4 != N % 4; j++) {
            matrix[i][j] -= matrix[i][k] * matrix[k][j];
        }
        matrix[i][k] = 0;
    }
    pthread_mutex_lock(&lock);
    done++;
    pthread_mutex_unlock(&lock);
    return NULL;
}

void pthread_row_sse_pool() {
    pthread_mutex_init(&lock, NULL);
    threadpool_t* threadpool = threadpool_create(NUM_THREADS, NUM_THREADS, 0);
    for (int k = 0; k < N; k++) {
        done = 0;
        for (int j = k + 1; j < N; j++) {
            matrix[k][j] /= matrix[k][k];
        }
        matrix[k][k] = 1;
        for (int i = 0; i < NUM_THREADS; i++) {
            parm[i] = threadParm_t{i, k};
            threadpool_add(threadpool, &pthread_row_sse_pool_function,
                           (void*)&parm[i], 0);
        }
        while (done != 4) {
            Sleep(0.01);
        }
    }
}

int main() {
    createMatrix();
    // long freq, head, tail;
    // QueryPerformanceFrequency((LARGE_INTEGER*)&freq);
    // QueryPerformanceCounter((LARGE_INTEGER*)&head);
    // LU();
    // QueryPerformanceCounter((LARGE_INTEGER*)&tail);
    // cout << "serial: " << (tail - head) * 1000.0 / freq << "ms" << endl;
    // createMatrix();
    // QueryPerformanceCounter((LARGE_INTEGER*)&head);
    pthread_col();
    // QueryPerformanceCounter((LARGE_INTEGER*)&tail);
    // cout << "thread col: " << (tail - head) * 1000.0 / freq << "ms" << endl;
    // createMatrix();
    // QueryPerformanceCounter((LARGE_INTEGER*)&head);
    // pthread_row_pool();
    // QueryPerformanceCounter((LARGE_INTEGER*)&tail);
    // cout << "thread col pool: " << (tail - head) * 1000.0 / freq << "ms"
    //      << endl;
    // createMatrix();
    // QueryPerformanceCounter((LARGE_INTEGER*)&head);
    // pthread_row();
    // QueryPerformanceCounter((LARGE_INTEGER*)&tail);
    // cout << "thread row: " << (tail - head) * 1000.0 / freq << "ms" << endl;
    // createMatrix();
    // QueryPerformanceCounter((LARGE_INTEGER*)&head);
    // pthread_row_sse();
    // QueryPerformanceCounter((LARGE_INTEGER*)&tail);
    // cout << "thread row sse: " << (tail - head) * 1000.0 / freq << "ms"
    // <<endl;
    pthread_row_sse_pool();
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < N; j++) {
            cout << matrix[i][j] << " ";
        }
        cout << endl;
    }
}
